"""
Coze WebSocket客户端实现
支持实时音频流上传和chat.update事件处理
"""

import asyncio
import json
import logging
import websockets
import base64
from typing import Optional, Dict, Any, Callable, List
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
import uuid

logger = logging.getLogger(__name__)


class WebSocketEventType(Enum):
    """WebSocket事件类型"""
    CHAT_UPDATE = "chat.update"
    INPUT_AUDIO_BUFFER_APPEND = "input_audio_buffer.append"
    INPUT_AUDIO_BUFFER_COMMIT = "input_audio_buffer.commit"
    INPUT_AUDIO_BUFFER_CLEAR = "input_audio_buffer.clear"
    RESPONSE_AUDIO_DELTA = "response.audio.delta"
    RESPONSE_AUDIO_DONE = "response.audio.done"
    RESPONSE_TEXT_DELTA = "response.text.delta"
    RESPONSE_TEXT_DONE = "response.text.done"
    SESSION_UPDATE = "session.update"
    ERROR = "error"


@dataclass
class WebSocketConfig:
    """WebSocket配置"""
    bot_id: str
    base_url: str = "wss://api.coze.cn"
    timeout: int = 30
    max_reconnect_attempts: int = 5
    reconnect_delay: float = 2.0
    audio_sample_rate: int = 24000
    audio_channels: int = 1
    audio_chunk_size: int = 1024


@dataclass
class AudioChunk:
    """音频数据块"""
    data: bytes
    timestamp: datetime
    sequence: int


@dataclass
class ChatUpdateData:
    """Chat更新数据"""
    chat_config: Optional[Dict[str, Any]] = None
    parameters: Optional[Dict[str, Any]] = None


class WebSocketEventHandler:
    """WebSocket事件处理器基类"""
    
    async def on_connect(self):
        """连接建立时调用"""
        pass
        
    async def on_disconnect(self):
        """连接断开时调用"""
        pass
        
    async def on_error(self, error: Dict[str, Any]):
        """错误事件处理"""
        logger.error(f"WebSocket错误: {error}")
        
    async def on_response_text_delta(self, delta: str):
        """文本响应增量处理"""
        pass
        
    async def on_response_text_done(self, text: str):
        """文本响应完成处理"""
        pass
        
    async def on_response_audio_delta(self, audio_data: bytes):
        """音频响应增量处理"""
        pass
        
    async def on_response_audio_done(self, audio_data: bytes):
        """音频响应完成处理"""
        pass
        
    async def on_session_update(self, session_data: Dict[str, Any]):
        """会话更新处理"""
        pass


class CozeWebSocketClient:
    """
    Coze WebSocket客户端
    支持实时音频流上传和对话交互
    """
    
    def __init__(self, config: WebSocketConfig, event_handler: Optional[WebSocketEventHandler] = None):
        """
        初始化WebSocket客户端
        
        Args:
            config: WebSocket配置
            event_handler: 事件处理器
        """
        self.config = config
        self.event_handler = event_handler or WebSocketEventHandler()
        self.websocket: Optional[websockets.WebSocketServerProtocol] = None
        self.is_connected = False
        self.reconnect_count = 0
        self.session_id: Optional[str] = None
        self.audio_sequence = 0
        self.response_text_buffer = ""
        self.response_audio_buffer = bytearray()
        
    async def connect(self) -> bool:
        """
        建立WebSocket连接
        
        Returns:
            连接是否成功
        """
        try:
            # 构建WebSocket URL
            url = f"{self.config.base_url}/v1/chat/websocket?bot_id={self.config.bot_id}"
            
            # 建立连接
            self.websocket = await websockets.connect(
                url,
                timeout=self.config.timeout,
                extra_headers={
                    "User-Agent": "EchokitBot-Vision/1.0"
                }
            )
            
            self.is_connected = True
            self.reconnect_count = 0
            self.session_id = str(uuid.uuid4())
            
            logger.info(f"WebSocket连接已建立: {url}")
            
            # 启动消息监听任务
            asyncio.create_task(self._message_listener())
            
            # 调用连接事件处理器
            await self.event_handler.on_connect()
            
            return True
            
        except Exception as e:
            logger.error(f"WebSocket连接失败: {e}")
            self.is_connected = False
            return False
            
    async def disconnect(self):
        """断开WebSocket连接"""
        if self.websocket and self.is_connected:
            try:
                await self.websocket.close()
                logger.info("WebSocket连接已断开")
            except Exception as e:
                logger.warning(f"断开WebSocket连接时出错: {e}")
            finally:
                self.websocket = None
                self.is_connected = False
                await self.event_handler.on_disconnect()
                
    async def _message_listener(self):
        """消息监听器"""
        try:
            async for message in self.websocket:
                await self._handle_message(message)
        except websockets.exceptions.ConnectionClosed:
            logger.warning("WebSocket连接已关闭")
            self.is_connected = False
            await self._handle_reconnect()
        except Exception as e:
            logger.error(f"消息监听器错误: {e}")
            self.is_connected = False
            
    async def _handle_message(self, message: str):
        """
        处理接收到的消息
        
        Args:
            message: 接收到的消息
        """
        try:
            data = json.loads(message)
            event_type = data.get("type")
            event_data = data.get("data", {})
            
            if event_type == WebSocketEventType.RESPONSE_TEXT_DELTA.value:
                delta = event_data.get("delta", "")
                self.response_text_buffer += delta
                await self.event_handler.on_response_text_delta(delta)
                
            elif event_type == WebSocketEventType.RESPONSE_TEXT_DONE.value:
                await self.event_handler.on_response_text_done(self.response_text_buffer)
                self.response_text_buffer = ""
                
            elif event_type == WebSocketEventType.RESPONSE_AUDIO_DELTA.value:
                audio_b64 = event_data.get("delta", "")
                if audio_b64:
                    audio_data = base64.b64decode(audio_b64)
                    self.response_audio_buffer.extend(audio_data)
                    await self.event_handler.on_response_audio_delta(audio_data)
                    
            elif event_type == WebSocketEventType.RESPONSE_AUDIO_DONE.value:
                await self.event_handler.on_response_audio_done(bytes(self.response_audio_buffer))
                self.response_audio_buffer.clear()
                
            elif event_type == WebSocketEventType.SESSION_UPDATE.value:
                await self.event_handler.on_session_update(event_data)
                
            elif event_type == WebSocketEventType.ERROR.value:
                await self.event_handler.on_error(event_data)
                
            else:
                logger.debug(f"未处理的事件类型: {event_type}")
                
        except json.JSONDecodeError as e:
            logger.error(f"消息解析失败: {e}")
        except Exception as e:
            logger.error(f"消息处理失败: {e}")
            
    async def _handle_reconnect(self):
        """处理重连逻辑"""
        if self.reconnect_count < self.config.max_reconnect_attempts:
            self.reconnect_count += 1
            delay = self.config.reconnect_delay * (2 ** (self.reconnect_count - 1))
            
            logger.info(f"尝试重连 ({self.reconnect_count}/{self.config.max_reconnect_attempts})，等待 {delay} 秒")
            await asyncio.sleep(delay)
            
            if await self.connect():
                logger.info("重连成功")
            else:
                await self._handle_reconnect()
        else:
            logger.error("达到最大重连次数，停止重连")
            
    async def _send_event(self, event_type: str, data: Dict[str, Any]):
        """
        发送事件消息
        
        Args:
            event_type: 事件类型
            data: 事件数据
        """
        if not self.is_connected or not self.websocket:
            raise RuntimeError("WebSocket未连接")
            
        message = {
            "type": event_type,
            "data": data
        }
        
        try:
            await self.websocket.send(json.dumps(message))
            logger.debug(f"发送事件: {event_type}")
        except Exception as e:
            logger.error(f"发送事件失败: {e}")
            raise
            
    async def send_image_parameter(self, image_file_id: str):
        """
        发送图像参数
        
        Args:
            image_file_id: 图像文件ID
        """
        chat_config = {
            "parameters": {
                "image": json.dumps({"file_id": image_file_id})
            }
        }
        
        await self._send_event(
            WebSocketEventType.CHAT_UPDATE.value,
            {"chat_config": chat_config}
        )
        
        logger.info(f"图像参数已发送: {image_file_id}")
        
    async def input_audio_buffer_append(self, audio_data: bytes):
        """
        追加音频数据到输入缓冲区
        
        Args:
            audio_data: 音频数据
        """
        # 将音频数据编码为base64
        audio_b64 = base64.b64encode(audio_data).decode('utf-8')
        
        await self._send_event(
            WebSocketEventType.INPUT_AUDIO_BUFFER_APPEND.value,
            {"delta": audio_b64}
        )
        
        self.audio_sequence += 1
        logger.debug(f"音频数据已追加: {len(audio_data)} 字节")
        
    async def input_audio_buffer_commit(self):
        """提交音频输入缓冲区"""
        await self._send_event(
            WebSocketEventType.INPUT_AUDIO_BUFFER_COMMIT.value,
            {}
        )
        logger.debug("音频输入缓冲区已提交")
        
    async def input_audio_buffer_clear(self):
        """清空音频输入缓冲区"""
        await self._send_event(
            WebSocketEventType.INPUT_AUDIO_BUFFER_CLEAR.value,
            {}
        )
        self.audio_sequence = 0
        logger.debug("音频输入缓冲区已清空")
        
    async def stream_audio(self, audio_data: bytes, chunk_size: Optional[int] = None):
        """
        流式上传音频数据
        
        Args:
            audio_data: 完整的音频数据
            chunk_size: 数据块大小
        """
        if chunk_size is None:
            chunk_size = self.config.audio_chunk_size
            
        # 清空缓冲区
        await self.input_audio_buffer_clear()
        
        # 分块发送音频数据
        for i in range(0, len(audio_data), chunk_size):
            chunk = audio_data[i:i + chunk_size]
            await self.input_audio_buffer_append(chunk)
            
            # 模拟实时音频流的延迟
            delay = len(chunk) * 1.0 / self.config.audio_sample_rate / self.config.audio_channels
            await asyncio.sleep(delay)
            
        # 提交音频数据
        await self.input_audio_buffer_commit()
        
        logger.info(f"音频流上传完成: {len(audio_data)} 字节")
        
    async def update_chat_config(self, chat_config: Dict[str, Any]):
        """
        更新对话配置
        
        Args:
            chat_config: 对话配置
        """
        await self._send_event(
            WebSocketEventType.CHAT_UPDATE.value,
            {"chat_config": chat_config}
        )
        
        logger.info("对话配置已更新")
        
    async def send_multimodal_query(
        self, 
        image_file_id: str, 
        audio_data: Optional[bytes] = None,
        parameters: Optional[Dict[str, Any]] = None
    ):
        """
        发送多模态查询
        
        Args:
            image_file_id: 图像文件ID
            audio_data: 音频数据（可选）
            parameters: 额外参数（可选）
        """
        # 发送图像参数
        await self.send_image_parameter(image_file_id)
        
        # 更新额外参数
        if parameters:
            chat_config = {"parameters": parameters}
            await self.update_chat_config(chat_config)
            
        # 发送音频数据
        if audio_data:
            await self.stream_audio(audio_data)
            
        logger.info("多模态查询已发送")


def split_bytes_by_length(data: bytes, chunk_size: int) -> List[bytes]:
    """
    按指定长度分割字节数据
    
    Args:
        data: 字节数据
        chunk_size: 块大小
        
    Returns:
        分割后的数据块列表
    """
    return [data[i:i + chunk_size] for i in range(0, len(data), chunk_size)]


class VisionWebSocketEventHandler(WebSocketEventHandler):
    """视觉交互专用的WebSocket事件处理器"""
    
    def __init__(self, response_callback: Optional[Callable[[str], None]] = None):
        """
        初始化事件处理器
        
        Args:
            response_callback: 响应回调函数
        """
        self.response_callback = response_callback
        self.current_response = ""
        
    async def on_connect(self):
        """连接建立时调用"""
        logger.info("视觉交互WebSocket连接已建立")
        
    async def on_disconnect(self):
        """连接断开时调用"""
        logger.info("视觉交互WebSocket连接已断开")
        
    async def on_response_text_delta(self, delta: str):
        """文本响应增量处理"""
        self.current_response += delta
        logger.debug(f"接收到文本增量: {delta}")
        
    async def on_response_text_done(self, text: str):
        """文本响应完成处理"""
        logger.info(f"接收到完整文本响应: {text}")
        
        if self.response_callback:
            self.response_callback(text)
            
        self.current_response = ""
        
    async def on_response_audio_delta(self, audio_data: bytes):
        """音频响应增量处理"""
        logger.debug(f"接收到音频增量: {len(audio_data)} 字节")
        
    async def on_response_audio_done(self, audio_data: bytes):
        """音频响应完成处理"""
        logger.info(f"接收到完整音频响应: {len(audio_data)} 字节")
        
    async def on_error(self, error: Dict[str, Any]):
        """错误事件处理"""
        logger.error(f"视觉交互WebSocket错误: {error}")